/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.facebook.presto.adaptive;

import com.facebook.presto.Session;
import com.facebook.presto.adaptive.selection.UnionIndexSelector;
import com.facebook.presto.metadata.Metadata;
import com.facebook.presto.metadata.MetadataManager;
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.ConnectorId;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.ConnectorTableHandle;
import com.facebook.presto.spi.ConnectorTableLayoutHandle;
import com.facebook.presto.spi.SchemaTableName;
import com.facebook.presto.spi.TableHandle;
import com.facebook.presto.spi.adaptive.CatalogEngine;
import com.facebook.presto.spi.adaptive.ConnectorAdaptiveOptimizer;
import com.facebook.presto.spi.adaptive.OptimizerType;
import com.facebook.presto.spi.adaptive.metadata.presto.AdaptivePrestoObject;
import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
import com.facebook.presto.spi.plan.FilterNode;
import com.facebook.presto.spi.plan.PlanNode;
import com.facebook.presto.spi.plan.PlanNodeIdAllocator;
import com.facebook.presto.spi.plan.TableScanNode;
import com.facebook.presto.spi.predicate.Domain;
import com.facebook.presto.spi.predicate.Range;
import com.facebook.presto.spi.predicate.SortedRangeSet;
import com.facebook.presto.spi.predicate.TupleDomain;
import com.facebook.presto.spi.relation.RowExpression;
import com.facebook.presto.spi.relation.VariableReferenceExpression;
import com.facebook.presto.spi.type.Type;
import com.facebook.presto.sql.TestingRowExpressionTranslator;
import com.facebook.presto.sql.planner.TypeProvider;
import com.facebook.presto.sql.planner.iterative.rule.test.PlanBuilder;
import com.facebook.presto.testing.TestingConnectorSession;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.testng.annotations.Test;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static com.facebook.presto.spi.type.BigintType.BIGINT;
import static com.facebook.presto.sql.planner.iterative.rule.test.PlanBuilder.expression;
import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
import static org.testng.Assert.assertNotEquals;
import static org.testng.AssertJUnit.assertEquals;

public class RBOAdaptiveSelector
{
    private static final Metadata METADATA = MetadataManager.createTestMetadataManager();
    private static final String CATALOG_NAME = "adptive_catalog";
    private static final String SCHEMA_NAME = "adptive_schema";
    private static final String TABLE_NAME = "adptive_table";
    private static final String CONNECTOR_ID = new ConnectorId(CATALOG_NAME).toString();
    private static final Session ADAPTIVE_SESSION = testSessionBuilder()
            .setCatalog(CATALOG_NAME)
            .setSchema(SCHEMA_NAME)
            .build();
    private static final ConnectorSession CONNECTOR_SESSION = new TestingConnectorSession(ImmutableList.of());
    private static final PlanNodeIdAllocator PLAN_NODE_ID_ALLOCATOR = new PlanNodeIdAllocator();
    private static final PlanBuilder PLAN_BUILDER = new PlanBuilder(ADAPTIVE_SESSION, PLAN_NODE_ID_ALLOCATOR, METADATA);

    private static final PlanNode NEW_NODE = PLAN_BUILDER.filter(new VariableReferenceExpression("test", BIGINT), null);

    private final TestingRowExpressionTranslator sqlToRowExpressionTranslator;

    public RBOAdaptiveSelector()
    {
        this.sqlToRowExpressionTranslator = new TestingRowExpressionTranslator(METADATA);
    }

    /**
     * sql:      select c1,c2 from test where c1 = 10;
     * index:    c1|c2
     * result:   satisfy unionIndexOptimize
     */
    @Test
    public void testUnionIndex1()
    {
        String expression = "c1 = 10";
        TypeProvider typeProvider = TypeProvider.copyOf(ImmutableMap.of("c1", BIGINT, "c2", BIGINT, "c3", BIGINT));
        RowExpression rowExpression = sqlToRowExpressionTranslator.translateAndOptimize(expression(expression), typeProvider);

        ColumnHandle column1 = new TestColumnHandle("c1");
        ColumnHandle column2 = new TestColumnHandle("c2");
        ColumnHandle column3 = new TestColumnHandle("c3");

        Domain domain1 = Domain.create(SortedRangeSet.copyOf(BIGINT, ImmutableList.of(Range.equal(BIGINT, 10L))), false);
        Domain domain2 = Domain.all(BIGINT);
        Domain domain3 = Domain.all(BIGINT);

        Map<ColumnHandle, Domain> columnDomains = ImmutableMap.of(column1, domain1, column2, domain2, column3, domain3);

        List<String> indexList = ImmutableList.of("c1|c2");
        List<String> outputColumns = ImmutableList.of("c1", "c2");

        TableScanNode tableScanNode = TableScan(SCHEMA_NAME, TABLE_NAME, BIGINT, columnDomains, outputColumns);

        FilterNode filterNode = filter(tableScanNode, rowExpression);

        UnionIndexSelector unionIndexSelector = new UnionIndexSelector();

        PlanNode newNode = unionIndexSelector.testFilterAdaptiveSelect(filterNode, indexList, PLAN_NODE_ID_ALLOCATOR, ImmutableList.of(new TestAdaptiveOptimizer(NEW_NODE)));

        assertNotEquals(filterNode, newNode);
    }

    /**
     * sql:      select c1 from test where c1 = 10;
     * index:    c2
     * result:   not satisfy unionIndexOptimize
     */
    @Test
    public void testUnionIndex2()
    {
        String expression = "c1 = 10";
        TypeProvider typeProvider = TypeProvider.copyOf(ImmutableMap.of("c1", BIGINT, "c2", BIGINT, "c3", BIGINT));
        RowExpression rowExpression = sqlToRowExpressionTranslator.translateAndOptimize(expression(expression), typeProvider);

        ColumnHandle column1 = new TestColumnHandle("c1");
        ColumnHandle column2 = new TestColumnHandle("c2");
        ColumnHandle column3 = new TestColumnHandle("c3");

        Domain domain1 = Domain.create(SortedRangeSet.copyOf(BIGINT, ImmutableList.of(Range.equal(BIGINT, 10L))), false);
        Domain domain2 = Domain.all(BIGINT);
        Domain domain3 = Domain.all(BIGINT);

        Map<ColumnHandle, Domain> columnDomains = ImmutableMap.of(column1, domain1, column2, domain2, column3, domain3);

        List<String> indexList = ImmutableList.of("c2");
        List<String> outputColumns = ImmutableList.of("c1");

        TableScanNode tableScanNode = TableScan(SCHEMA_NAME, TABLE_NAME, BIGINT, columnDomains, outputColumns);

        FilterNode filterNode = filter(tableScanNode, rowExpression);

        UnionIndexSelector unionIndexSelector = new UnionIndexSelector();

        PlanNode newNode = unionIndexSelector.testFilterAdaptiveSelect(filterNode, indexList, PLAN_NODE_ID_ALLOCATOR, ImmutableList.of(new TestAdaptiveOptimizer(NEW_NODE)));

        assertEquals(filterNode, newNode);
    }

    /**
     * sql:      select c1 from test where c1 = 10;
     * index:    c2,c3
     * result:   not satisfy unionIndexOptimize
     */
    @Test
    public void testUnionIndex3()
    {
        String expression = "c1 = 10";
        TypeProvider typeProvider = TypeProvider.copyOf(ImmutableMap.of("c1", BIGINT, "c2", BIGINT, "c3", BIGINT));
        RowExpression rowExpression = sqlToRowExpressionTranslator.translateAndOptimize(expression(expression), typeProvider);

        ColumnHandle column1 = new TestColumnHandle("c1");
        ColumnHandle column2 = new TestColumnHandle("c2");
        ColumnHandle column3 = new TestColumnHandle("c3");

        Domain domain1 = Domain.create(SortedRangeSet.copyOf(BIGINT, ImmutableList.of(Range.equal(BIGINT, 10L))), false);
        Domain domain2 = Domain.all(BIGINT);
        Domain domain3 = Domain.all(BIGINT);

        Map<ColumnHandle, Domain> columnDomains = ImmutableMap.of(column1, domain1, column2, domain2, column3, domain3);

        List<String> indexList = ImmutableList.of("c2", "c3");
        List<String> outputColumns = ImmutableList.of("c1");

        TableScanNode tableScanNode = TableScan(SCHEMA_NAME, TABLE_NAME, BIGINT, columnDomains, outputColumns);

        FilterNode filterNode = filter(tableScanNode, rowExpression);

        UnionIndexSelector unionIndexSelector = new UnionIndexSelector();

        PlanNode newNode = unionIndexSelector.testFilterAdaptiveSelect(filterNode, indexList, PLAN_NODE_ID_ALLOCATOR, ImmutableList.of(new TestAdaptiveOptimizer(NEW_NODE)));

        assertEquals(filterNode, newNode);
    }

    /**
     * sql:      select c1 from test;
     * index:    c2,c3
     * result:   not satisfy unionIndexOptimize
     */
    @Test
    public void testUnionIndex4()
    {
        String expression = "c1 = 10";
        TypeProvider typeProvider = TypeProvider.copyOf(ImmutableMap.of("c1", BIGINT, "c2", BIGINT, "c3", BIGINT));

        ColumnHandle column1 = new TestColumnHandle("c1");
        ColumnHandle column2 = new TestColumnHandle("c2");
        ColumnHandle column3 = new TestColumnHandle("c3");

        Domain domain1 = Domain.all(BIGINT);
        Domain domain2 = Domain.all(BIGINT);
        Domain domain3 = Domain.all(BIGINT);

        Map<ColumnHandle, Domain> columnDomains = ImmutableMap.of(column1, domain1, column2, domain2, column3, domain3);

        List<String> indexList = ImmutableList.of("c2", "c3");
        List<String> outputColumns = ImmutableList.of("c1");

        TableScanNode tableScanNode = TableScan(SCHEMA_NAME, TABLE_NAME, BIGINT, columnDomains, outputColumns);

        UnionIndexSelector unionIndexSelector = new UnionIndexSelector();

        PlanNode newNode = unionIndexSelector.testScanAdaptiveSelect(tableScanNode, indexList, PLAN_NODE_ID_ALLOCATOR, ImmutableList.of(new TestAdaptiveOptimizer(NEW_NODE)));

        assertEquals(tableScanNode, newNode);
    }

    /**
     * sql:      select c1 from test;
     * index:    c1,c2
     * result:   satisfy unionIndexOptimize
     */
    @Test
    public void testUnionIndex5()
    {
        String expression = "c1 = 10";
        TypeProvider typeProvider = TypeProvider.copyOf(ImmutableMap.of("c1", BIGINT, "c2", BIGINT, "c3", BIGINT));

        ColumnHandle column1 = new TestColumnHandle("c1");
        ColumnHandle column2 = new TestColumnHandle("c2");
        ColumnHandle column3 = new TestColumnHandle("c3");

        Domain domain1 = Domain.all(BIGINT);
        Domain domain2 = Domain.all(BIGINT);
        Domain domain3 = Domain.all(BIGINT);

        Map<ColumnHandle, Domain> columnDomains = ImmutableMap.of(column1, domain1, column2, domain2, column3, domain3);

        List<String> indexList = ImmutableList.of("c1", "c2");
        List<String> outputColumns = ImmutableList.of("c1");

        TableScanNode tableScanNode = TableScan(SCHEMA_NAME, TABLE_NAME, BIGINT, columnDomains, outputColumns);

        UnionIndexSelector unionIndexSelector = new UnionIndexSelector();

        PlanNode newNode = unionIndexSelector.testScanAdaptiveSelect(tableScanNode, indexList, PLAN_NODE_ID_ALLOCATOR, ImmutableList.of(new TestAdaptiveOptimizer(NEW_NODE)));

        assertNotEquals(tableScanNode, newNode);
    }

    private TableScanNode TableScan(String schema, String table, Type type, Map<ColumnHandle, Domain> columnDomains, List<String> ouputColumns)
    {
        TupleDomain<ColumnHandle> tupleDomain = TupleDomain.withColumnDomains(columnDomains);

        ConnectorTableHandle testConnectorTableHandle = new TestConnectorTableHandle(SCHEMA_NAME, TABLE_NAME);
        ConnectorTableLayoutHandle testTableLayoutHandle = new TestTableLayoutHandle(tupleDomain);
        TableHandle tableHandle = new TableHandle(new ConnectorId(CATALOG_NAME), testConnectorTableHandle, new ConnectorTransactionHandle() {}, Optional.of(testTableLayoutHandle));

        List<VariableReferenceExpression> outputVariables = new ArrayList<>();
        Map<VariableReferenceExpression, ColumnHandle> assignments = new HashMap<>();

        for (Map.Entry<ColumnHandle, Domain> columnDomain : columnDomains.entrySet()) {
            ColumnHandle columnHandle = columnDomain.getKey();
            String columnName = columnHandle.getColumnName();

            VariableReferenceExpression vRE = newVariable(columnName, type);

            if (ouputColumns.contains(columnName)) {
                outputVariables.add(vRE);
            }
            assignments.put(vRE, columnHandle);
        }

        return PLAN_BUILDER.tableScan(
                tableHandle,
                outputVariables,
                assignments);
    }

    private FilterNode filter(PlanNode source, RowExpression predicate)
    {
        return PLAN_BUILDER.filter(predicate, source);
    }

    private static VariableReferenceExpression newVariable(String name, Type type)
    {
        return new VariableReferenceExpression(name, type);
    }
}

class TestConnectorTableHandle
        implements ConnectorTableHandle
{
    private SchemaTableName schemaTableName;

    public TestConnectorTableHandle(String schema, String table)
    {
        this.schemaTableName = new SchemaTableName(schema, table);
    }

    @Override
    public SchemaTableName getSchemaTableName()
    {
        return this.schemaTableName;
    }
}

class TestTableLayoutHandle
        implements ConnectorTableLayoutHandle
{
    TupleDomain<ColumnHandle> tupleDomain;

    public TestTableLayoutHandle(TupleDomain<ColumnHandle> tupleDomain)
    {
        this.tupleDomain = tupleDomain;
    }

    public TupleDomain<ColumnHandle> getTupleDomain()
    {
        return this.tupleDomain;
    }
}

class TestColumnHandle
        implements ColumnHandle
{
    private String columnName;

    public TestColumnHandle(String key, Domain value)
    {
    }

    public TestColumnHandle(String columnName)
    {
        this.columnName = columnName;
    }

    @Override
    public String getColumnName()
    {
        return columnName;
    }
}

class TestAdaptiveOptimizer
        implements ConnectorAdaptiveOptimizer
{

    private PlanNode newNode;

    public TestAdaptiveOptimizer(PlanNode newNode)
    {
        this.newNode = newNode;
    }

    @Override
    public CatalogEngine getConnectorCatalog()
    {
        return CatalogEngine.TEST;
    }

    @Override
    public PlanNode adaptiveOptimize(PlanNode node, TableHandle tableHandle, PlanNodeIdAllocator idAllocator, AdaptivePrestoObject adaptiveObject)
    {
        return newNode;
    }

    @Override
    public OptimizerType getOptimizerType()
    {
        return OptimizerType.TEST;
    }
}
